package main

import (
	"context"
	"testing"

	"github.com/pingcap/kvproto/pkg/kvrpcpb"
	"github.com/stretchr/testify/assert"
	"github.com/tikv/client-go/v2/rawkv"
)

func cleanKeys(client *rawkv.Client, t *testing.T, key ...[]byte) {
	for _, k := range key {
		req := buildRedisRequest("DEL", k)
		if _, err := client.RedisExec(context.TODO(), req); err != nil {
			t.Fatalf("del error: %v", err)
		}
	}
}

func buildRedisRequest(cmd string, key []byte, args ...[]byte) *kvrpcpb.RedisRequest {
	return &kvrpcpb.RedisRequest{
		Cmd:  cmd,
		Key:  key,
		Args: args,
	}
}

func assertOK(t *testing.T, resp *kvrpcpb.RedisResponse, err error) {
	if err != nil {
		t.Fatalf("redis cmd error: %v", err)
	}
	assert.NotNil(t, resp.GetSimple())
	assert.Equal(t, "OK", resp.GetSimple().GetMessage())
}

func assertBulkMsg(t *testing.T, resp *kvrpcpb.RedisResponse, msg []byte, err error) {
	if err != nil {
		t.Fatalf("redis cmd error: %v", err)
	}
	assert.NotNil(t, resp.GetBulk())
	assert.Equal(t, msg, resp.GetBulk().GetValue())
}

func assertBulkNil(t *testing.T, resp *kvrpcpb.RedisResponse, err error) {
	if err != nil {
		t.Fatalf("redis cmd error: %v", err)
	}
	assert.NotNil(t, resp.GetBulk())
	assert.True(t, resp.GetBulk().GetIsNil())
}

func assertInt(t *testing.T, resp *kvrpcpb.RedisResponse, i int64, err error) {
	if err != nil {
		t.Fatalf("redis cmd error: %v", err)
	}
	assert.NotNil(t, resp.GetInteger())
	assert.Equal(t, i, resp.GetInteger().GetValue())
}

func assertNotError(t *testing.T, resp *kvrpcpb.RedisResponse, err error) {
	if err != nil {
		t.Fatalf("redis cmd error: %s", err.Error())
	} else if resp.RegionError != nil {
		t.Fatalf("region error: %s", resp.RegionError)
	}
	assert.Nil(t, resp.GetError())
}

func assertHasError(t *testing.T, resp *kvrpcpb.RedisResponse, err error) {
	if err != nil {
		t.Fatalf("redis cmd error: %s", err.Error())
	} else if resp.RegionError != nil {
		t.Fatalf("region error: %s", resp.RegionError)
	}
	assert.NotNil(t, resp.GetError())
}

func assertArrayBulk(t *testing.T, resp *kvrpcpb.RedisResponse, elements [][]byte, err error) {
	if err != nil {
		t.Fatalf("redis cmd error: %v", err)
	}
	assert.NotNil(t, resp.GetArray())
	assert.Equal(t, len(elements), len(resp.GetArray().GetElements()))

	for i, value := range resp.GetArray().GetElements() {
		// 只判断元素为bulk的情况
		assert.NotNil(t, value.GetBulk())
		if elements[i] == nil {
			assert.True(t, value.GetBulk().GetIsNil())
		} else {
			assert.Equal(t, elements[i], value.GetBulk().GetValue())
		}
	}
}

func assertArrayInt(t *testing.T, resp *kvrpcpb.RedisResponse, elements []int, err error) {
	if err != nil {
		t.Fatalf("redis cmd error: %v", err)
	}
	assert.NotNil(t, resp.GetArray())
	assert.Equal(t, len(elements), len(resp.GetArray().GetElements()))

	for i, value := range resp.GetArray().GetElements() {
		assert.NotNil(t, value.GetInteger())
		assert.Equal(t, int64(elements[i]), value.GetInteger().GetValue())
	}
}

func assertArrayNil(t *testing.T, resp *kvrpcpb.RedisResponse, err error) {
	if err != nil {
		t.Fatalf("redis cmd error: %v", err)
	}
	assert.NotNil(t, resp.GetArray())

	for _, value := range resp.GetArray().GetElements() {
		assert.NotNil(t, value.GetBulk())
		assert.True(t, value.GetBulk().GetIsNil())
	}
}

func assertScanResponse(t *testing.T, resp *kvrpcpb.RedisResponse, err error) (int, string) {
	if err != nil {
		t.Fatalf("redis cmd error: %v", err)
	}

	array := resp.GetArray().GetElements()
	assert.Equal(t, 2, len(array))

	cursorMsg := array[0].GetBulk()
	assert.NotNil(t, cursorMsg)
	assert.False(t, cursorMsg.GetIsNil())

	elementMsg := array[1].GetArray()
	assert.NotNil(t, elementMsg)

	count := 0
	if !elementMsg.GetIsNull() {
		count = len(elementMsg.GetElements())
	}

	cursor := string(cursorMsg.GetValue())
	return count, cursor
}

func distinctCount(resp []*kvrpcpb.RedisResponse) int {
	set := make(map[string]struct{})
	for _, r := range resp {
		set[r.GetBulk().String()] = struct{}{}
	}
	return len(set)
}